import logging
import json
import os
import re
from typing import Optional
from core.testAgent.LoggerManager import LoggerManager
from core.testAgent.CKGRetriever import CKGRetriever
from core.testAgent.EnvironmentService import EnvironmentService
from core.testAgent.Planner import PlannerState
from core.testAgent.Planner import graph
from core.testAgent.Config import *
from core.testAgent.Utils import add_record, add_test_case
from core.testAgent.ProgressTracker import progress_tracker

def get_graph_retriever(url, username, password):
    """
    Create a new instance of CKGRetriever with the provided URL, username, and password.
    """
    return CKGRetriever(url, username, password)

def list_methods(url, username, password, limit=50, offset=0, focal_only=True):
    """
    Fetch a concise list of methods from Neo4j so the VS Code command can display them.
    """
    graph_retriever = get_graph_retriever(url, username, password)
    where_clause = """
        WHERE m.modifiers CONTAINS "public"
          AND NOT m.modifiers CONTAINS "abstract"
          AND m.type = "METHOD"
    """ if focal_only else ""

    query = f"""
        MATCH (m:Method)
        {where_clause}
        RETURN id(m) AS method_id,
               m.name AS name,
               m.class_name AS class_name,
               m.package_name AS package_name,
               m.signature AS signature
        ORDER BY m.name ASC
        SKIP $offset
        LIMIT $limit
    """

    methods = []
    with graph_retriever.driver.session() as session:
        records = session.run(query, offset=offset, limit=limit)
        for record in records:
            methods.append({
                "id": int(record["method_id"]),
                "name": record.get("name") or "",
                "class_name": record.get("class_name") or "",
                "package_name": record.get("package_name") or "",
                "signature": record.get("signature") or "",
            })
    return methods

def generate_test_case(method_id, url, username, password, limit, job_id: Optional[str] = None,
                       allow_human_assistance: bool = True, project_path: Optional[str] = None):
    progress_tracker.start(job_id, {"methodId": method_id})
    config = {"recursion_limit": limit}
    graph_retriever = get_graph_retriever(url, username, password)
    method = graph_retriever.search_method_by_id(method_id)
    if isinstance(method, str):
        error_message = f"Method lookup failed for id {method_id}: {method}"
        progress_tracker.fail(job_id, error_message)
        return {
            "result": "Failed",
            "message": error_message
        }
    progress_tracker.record_message(job_id, f"Fetched method {getattr(method, 'full_qualified_name', method_id)}")
    envServer = EnvironmentService()
    envServer.set_job_id(job_id)
    normalized_project_path = _normalize_fs_path(project_path)
    base_path = _resolve_base_path(method, normalized_project_path)
    if not base_path or not os.path.isdir(base_path):
        error_message = (
            f"Unable to resolve project base path for method {getattr(method, 'full_qualified_name', method_id)}. "
            "Please ensure the knowledge graph contains absolute paths or provide a valid project path."
        )
        progress_tracker.fail(job_id, error_message)
        return {
            "result": "Failed",
            "message": error_message
        }
    envServer.set_base_path(base_path)
    envServer.set_human_assistance_enabled(allow_human_assistance)
    # 使用 element_id (字符串) 设置 focal_method_id
    graph_retriever.change_focal_method_id(method.id)
    inject_dir = _resolve_inject_dir(method, base_path, normalized_project_path)
    if not inject_dir:
        error_message = (
            f"Unable to resolve injection directory for method {getattr(method, 'full_qualified_name', method_id)}. "
            "Please ensure either the knowledge graph stores absolute paths or specify a project path."
        )
        progress_tracker.fail(job_id, error_message)
        return {
            "result": "Failed",
            "message": error_message
        }
    envServer.set_inject_dir(inject_dir)
    print("================= TEST GENERATE START =================\n")
    print(f"package_name: {method.package_name}\n")
    print(f"method_name: {method.name}\n")
    print(f"method_signature: {method.signature}\n")
    print(f"full_qualified_name: {method.full_qualified_name}\n")
    print(f"absolute_path: {method.absolute_path}\n\n")
    # 使用原始的整数 method_id (参数传入的)
    state = PlannerState(envServer=envServer,
                            package_name=method.package_name,
                            method_id=method_id,  # 使用传入的整数 ID
                            job_id=job_id,
                            allow_human_assistance=allow_human_assistance,
                            method_code=method.content,
                            method_signature=method.signature,
                            class_name=method.class_name,
                            full_method_name=method.full_qualified_name,
                            start_line=method.start_line,
                            end_line=method.end_line)
    success = False
    try:
        for event in graph.stream(state, config, subgraphs=True):
            print("========================== Graph:" + str(event[0]) + " =====================\n")
            node_name = str(list(event[1].keys())[0])
            print("Current Node: " + node_name + "\n")
            progress_tracker.record_node(job_id, node_name)
            for value in event[1].values():
                if 'messages' in value.keys():
                    for m in value["messages"]:
                        print(m.pretty_repr() + '\n\n')
                else:
                    print('\n' + str(value) + '\n')
            event = event[1]
            EndPoint = 'testCaseAcceptor'
            if EndPoint in event.keys():
                print("===================== TEST GENERATE END =====================\n\n")
                result = []
                test_code = []
                find_bugs = []
                covered_lines = set()
                total_lines = set()
                mutation_info = {}
                line_coverage_rates = []
                branch_coverage_rates = []
                mutation_scores = []
                mutants_list = []
                test_point = []
                test_report = []
                inner_index = 0
                for test_case in event[EndPoint]['test_cases']:
                    print(f"test case: {(test_case['test_case'])}\n")
                    inner_index += 1
                    print(f"===================== Test Report {inner_index} =====================\n")
                    print(f"test result: {(test_case['test_result'])}\n")
                    print(f"find bugs: {(test_case['find_bug'])}\n")
                    print(f"test case: {(test_case['test_case'])}\n")
                    print(f"test point: {(test_case['test_point'])}\n")
                    print(f"test report: {(test_case['test_report'])}\n")
                    print(f"coverage report: {(test_case['coverage_report'])}\n")
                    print(f"mutation report: {(test_case['mutation_report'])}\n")
                    test_status = test_case['test_result']
                    if test_status in ("Compile Error", "Syntax Error"):
                        print(f"Skip persisting invalid test case due to result: {test_status}\n")
                        continue
                    result.append(test_status)
                    test_code.append(test_case['test_case'])
                    find_bugs.append(test_case['find_bug'])
                    test_point.append(test_case['test_point'])
                    test_report.append(test_case['test_report'])
                    if test_case['coverage_report']['result'] == 'Success':
                        line_coverage_rates.append(test_case['coverage_report']['output']['line_coverage'])
                        branch_coverage_rates.append(test_case['coverage_report']['output']['branch_coverage'])
                        total_lines.update(test_case['coverage_report']['output']['covered_lines'])
                        total_lines.update(test_case['coverage_report']['output']['missed_lines'])
                        covered_lines.update(test_case['coverage_report']['output']['covered_lines'])
                    else:
                        line_coverage_rates.append(0)
                        branch_coverage_rates.append(0)
                    if test_case['mutation_report']['result'] == 'Success':
                        mutation_scores.append(test_case['mutation_report']['output']['mutation_score'])
                        mutants_list.append(test_case['mutation_report']['output']['filtered_mutations'])
                        for item_mutation in test_case['mutation_report']['output']['filtered_mutations']:
                            key = str(item_mutation['Line']) + '_' + str(item_mutation['Mutator'])
                            if key in mutation_info.keys() and mutation_info[key] == 'KILLED':
                                continue
                            else:
                                mutation_info[key] = item_mutation['Status']
                    else:
                        mutation_scores.append(0)

                if not test_code:
                    progress_tracker.fail(job_id, "No compilable test case generated.")
                    return {
                        "result": "Failed",
                        "output": "No compilable test case generated."
                    }

                print(f"test result: {(result)}\n")

                cal_coverage = len(covered_lines) / len(total_lines) if len(total_lines) != 0 else 0
                max_line_rate = max(line_coverage_rates) if line_coverage_rates else 0
                final_coverage = cal_coverage if cal_coverage > max_line_rate / 100 else max_line_rate / 100
                print(f"final coverage: {final_coverage}\n")
                final_mutation_score = sum([1 for v in mutation_info.values() if v == 'KILLED']) / len(
                    mutation_info) if len(mutation_info) != 0 else 0
                print(f"final mutation score: {final_mutation_score}\n")

                saved_files = persist_test_files(
                    test_code,
                    inject_dir,
                    method.class_name or method.name,
                    method.name
                )
                success = True
                progress_tracker.finish(job_id, "success", "Test generation complete")
                return {
                    "result": "Success",
                    "output": {
                        "test_code": test_code,
                        "test_report": test_report,
                        "find_bugs": find_bugs,
                        "test_point": test_point,
                        "line_coverage_rate": line_coverage_rates,
                        "branch_coverage_rate": branch_coverage_rates,
                        "mutation_score": mutation_scores,
                        "final_coverage": final_coverage,
                        "final_mutation_score": final_mutation_score,
                        "mutation_info": mutation_info,
                        "saved_files": saved_files,
                        "saved_file": saved_files[0] if saved_files else None
                    }
                }
        progress_tracker.fail(job_id, "No test case generated.")
        return {
            "result": "Failed",
            "output": "No test case generated."
        }
    except Exception as exc:
        if not success:
            progress_tracker.fail(job_id, str(exc))
        raise


def persist_test_files(test_code_blocks, inject_dir, class_name, method_name):
    saved_files = []
    if not test_code_blocks or not inject_dir:
        return saved_files

    try:
        os.makedirs(inject_dir, exist_ok=True)
        sanitized_class = _sanitize_name(class_name) or "Generated"
        sanitized_method = _sanitize_name(method_name) or "Method"

        for idx, code in enumerate(test_code_blocks, start=1):
            base_name = f"{sanitized_class}_{sanitized_method}_{idx}_Test"
            base_name = _sanitize_name(base_name) or f"{sanitized_class}_{sanitized_method}_{idx}_Test"
            updated_code = _replace_public_class_name(code, base_name)

            file_name = f"{base_name}.java"
            file_path = os.path.join(inject_dir, file_name)
            with open(file_path, "w", encoding="utf-8") as file_handle:
                file_handle.write(updated_code)
            saved_files.append(file_path)
        return saved_files
    except Exception as exc:
        logging.error(f"Failed to persist test files: {exc}")
        return saved_files


def _sanitize_name(raw_name):
    if not raw_name:
        return ""
    sanitized = re.sub(r'[^0-9a-zA-Z_]+', '_', raw_name)
    return sanitized


def _replace_public_class_name(java_code: str, new_class_name: str) -> str:
    """
    Replace the first public class declaration with the provided class name to
    keep the filename and class name consistent.
    """
    if not java_code:
        return java_code

    pattern = re.compile(r'(public\s+(?:\w+\s+)*)class\s+(\w+)')

    def repl(match):
        prefix = match.group(1)
        return f"{prefix}class {new_class_name}"

    updated_code, count = pattern.subn(repl, java_code, count=1)
    if count == 0:
        logging.warning("No public class declaration found; class name could not be aligned with filename.")
        return java_code
    return updated_code


def _normalize_fs_path(path: Optional[str]) -> str:
    if not path:
        return ""
    normalized = path.replace("\\", os.sep).strip()
    if not normalized:
        return ""
    return os.path.abspath(normalized)


def _resolve_base_path(method, fallback_project_path: Optional[str]) -> str:
    absolute_path = getattr(method, "absolute_path", "") or ""
    normalized = absolute_path.replace("\\", "/").strip()
    candidate = ""
    if normalized:
        candidate = normalized.split("/src")[0]
    candidate = _normalize_fs_path(candidate)
    if candidate and os.path.isdir(candidate):
        return candidate
    if fallback_project_path and os.path.isdir(fallback_project_path):
        return fallback_project_path
    return candidate


def _resolve_inject_dir(method, base_path: str, fallback_project_path: Optional[str]) -> str:
    absolute_path = getattr(method, "absolute_path", "") or ""
    normalized = absolute_path.replace("\\", "/").strip()
    candidate = ""
    if normalized:
        replaced = normalized.replace("src/main", "src/test")
        candidate = os.path.dirname(replaced)
    candidate = _normalize_fs_path(candidate)
    if candidate:
        return candidate
    project_root = base_path or fallback_project_path
    if not project_root:
        return ""
    package_name = getattr(method, "package_name", "") or ""
    package_parts = [part for part in package_name.split(".") if part]
    fallback_dir = os.path.join(project_root, "src", "test", "java", *package_parts)
    return _normalize_fs_path(fallback_dir)


if __name__ == '__main__':
    # assert len(methods_list) < 201
    method_id = 44777
    ret = generate_test_case(method_id, "bolt://localhost:7687", "neo4j", "123456", 50, None, True, None)
